from sklearn.neighbors import KNeighborsClassifier

# X 中存储交通工具的参数，包括总长度（米）、时速（km/h）、重量（吨）、座位数量
X = [
    [96, 85, 120, 400],         # 普通火车
    [144, 92, 200, 600],
    [240, 87, 350, 1000],
    [360, 90, 495, 1300],
    [384, 91, 530, 1405],
    [240, 360, 490, 800],       # 高铁
    [360, 380, 750, 1200],
    [290, 380, 480, 960],
    [120, 320, 160, 400],
    [384, 340, 520, 1280],
    [33.4, 918, 77, 180],       # 飞机
    [33.6, 1120, 170.5, 185],
    [39.5, 785, 230, 240],
    [33.84, 940, 150, 195],
    [44.5, 920, 275, 275],
    [75.3, 1050, 575, 490]
]

# y 中存储类型，0表示普通火车，1表示高铁，2表示飞机
y = [0] * 5 + [1] * 5 + [2] * 6

# labels 中存储对应的交通工具名称
labels = ('普通火车', '高铁', '飞机')

# 创建并训练模型
knn = KNeighborsClassifier(n_neighbors=3, weights='distance')
knn.fit(X, y)

# 对未知数据进行分类
unKnown = [[300, 79, 320, 900], [36.7, 800, 190, 220]]
result = knn.predict(unKnown)
for para, index in zip(unKnown, result):
    print(para, labels[index], sep=':')
